#ifndef BRDF_GLSL
#define BRDF_GLSL

#include "Common.glsl"

mat3 MakeRotationMatrix(vec3 z)
{
    const vec3 ref = abs(dot(z, vec3(0, 1, 0))) > 0.99f ? vec3(0, 0, 1) : vec3(0, 1, 0);

    const vec3 x = normalize(cross(ref, z));
    const vec3 y = cross(z, x);

    return mat3(x, y, z);
}

vec3 SampleCosineLobe(in vec3 n, in vec2 r)
{
    vec2 randSample = max(vec2(0.00001f), r);

    const float phi = 2.0f * M_PI * randSample.y;

    const float cosTheta = sqrt(randSample.x);
    const float sinTheta = sqrt(1.0f - randSample.x);

    vec3 t = vec3(sinTheta * cos(phi), sinTheta * sin(phi), cosTheta);

    return normalize(MakeRotationMatrix(n) * t);
}



float D_GGX(in float dotNH, in float alpha)
{
    float a2 = alpha * alpha;
    float denom = (dotNH * dotNH) * (a2 - 1.0) + 1.0;

    return a2 / max(EPSILON, (M_PI * denom * denom));
}



float G1_Schlick_GGX(in float roughness, in float dotNV)
{
    float k = ((roughness + 1) * (roughness + 1)) / 8.0;

    return dotNV / max(EPSILON, (dotNV * (1 - k) + k));
}



float G_Schlick_GGX(in float dotNL, in float dotNV, in float roughness)
{
    return G1_Schlick_GGX(roughness, dotNL) * G1_Schlick_GGX(roughness, dotNV);
}



vec3 F_Schlick(in vec3 f0, in float dotVH)
{
    return f0 + (vec3(1.0) - f0) * (pow(1.0 - dotVH, 5.0));
}



vec3 EvaluateSpecularBrdf(in float roughness, in vec3 F, in float dotNH, in float dotNL, in float dotNV)
{
    float alpha = roughness * roughness;
    return (D_GGX(dotNH, alpha) * F * G_Schlick_GGX(dotNL, dotNV, roughness)) / max(EPSILON, (4.0 * dotNL * dotNV));
}



float PdfSpecularGGXLobe(in float alpha, in float dotNH, in float dotVH)
{
    return D_GGX(dotNH, alpha) * dotNH / max(EPSILON, (4.0 * dotVH));
}



float PdfCosineLobe(in float ndotl)
{
    return ndotl / M_PI;
}



vec3 EvaluateDiffuseBrdf(in vec3 diffuseColor)
{
    return diffuseColor / M_PI;
}



vec3 SampleSpecularGGXLobe(in vec3 n, in float alpha, in vec2 Xi)
{
    float phi       = 2.0 * M_PI * Xi.x;
    float cosTheta = sqrt((1.0 - Xi.y) / (1.0 + (alpha * alpha - 1.0) * Xi.y));
    float sinTheta = sqrt(1.0 - cosTheta * cosTheta);

    vec3 d;

    d.x = sinTheta * cos(phi);
    d.y = sinTheta * sin(phi);
    d.z = cosTheta;

    return normalize(MakeRotationMatrix(n) * d);
}



float PdfUberBrdf(in vec3 N, in float roughness, in vec3 Wo, in vec3 Wh, in vec3 Wi)
{
    float dotNL = max(dot(N, Wi), 0.0);
    float dotNV = max(dot(N, Wo), 0.0);
    float dotNH = max(dot(N, Wh), 0.0);
    float dotVH = max(dot(Wi, Wh), 0.0);

    float pd = PdfCosineLobe(dotNL);
    float ps = PdfSpecularGGXLobe(roughness * roughness, dotNH, dotVH);

    return mix(pd, ps, 0.5);
}



vec3 EvaluateUberBrdf(in vec3 diffuseColor, in float roughness, in vec3 N, in vec3 F0, in vec3 Wo, in vec3 Wh, in vec3 Wi)
{
    float dotNL = max(dot(N, Wi), 0.0);
    float dotNV = max(dot(N, Wo), 0.0);
    float dotNH = max(dot(N, Wh), 0.0);
    float dotVH = max(dot(Wi, Wh), 0.0);

    vec3 F = F_Schlick(F0, dotVH);
    vec3 specular = EvaluateSpecularBrdf(roughness, F, dotNH, dotNL, dotNV);
    vec3 diffuse  = EvaluateDiffuseBrdf(diffuseColor.xyz);

    return (vec3(1.0) - F) * diffuse + specular;
}



vec3 SampleUberBrdf(in vec3 diffuseColor, in vec3 F0, in vec3 N, in float roughness, in float metallic, in vec3 Wo, in RNG rng, out vec3 Wi, out float pdf)
{
    float alpha = roughness * roughness;

    vec3 Wh;

    vec3 randValue = NextVec3(rng);

    if (randValue.x < 0.5f)
    {
        Wh = SampleSpecularGGXLobe(N, alpha, randValue.yz);

        if (roughness < MIRROR_REFLECTIONS_ROUGHNESS_THRESHOLD)
            Wi = reflect(-Wo, N);
        else
            Wi = reflect(-Wo, Wh);
    }
    else
    {
        Wi = SampleCosineLobe(N, randValue.yz);
        Wh = normalize(Wo + Wi);
    }

    pdf = PdfUberBrdf(N, roughness, Wo, Wh, Wi);

    return EvaluateUberBrdf(diffuseColor, roughness, N, F0, Wo, Wh, Wi);
}



#endif